package com.ics.atable.chat.serializer;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.serializer.plain_text.PlainTextStateSerializer;
import com.alibaba.cloud.ai.graph.state.AgentStateFactory;
import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.PropertyAccessor;
import com.fasterxml.jackson.databind.MapperFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import org.springframework.ai.chat.messages.Message;

import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectOutput;
import java.nio.charset.StandardCharsets;

/**
 *
 * @author haohongbin
 */
public class ItemQueryStateSerializer extends PlainTextStateSerializer {

    protected final ObjectMapper objectMapper;

    public ItemQueryStateSerializer(AgentStateFactory<OverAllState> stateFactory) {
        this(stateFactory, null);
    }

    protected ItemQueryStateSerializer(AgentStateFactory<OverAllState> stateFactory,
                                       ObjectMapper customObjectMapper) {
        super(stateFactory);
        if (customObjectMapper != null) {
            this.objectMapper = customObjectMapper;
            return;
        }
        this.objectMapper = JsonMapper.builder()
                .configure(MapperFeature.AUTO_DETECT_GETTERS, false)
                .configure(MapperFeature.AUTO_DETECT_IS_GETTERS, false)
                .visibility(PropertyAccessor.FIELD, JsonAutoDetect.Visibility.ANY)
                .build();

        // register MessageDeserializer， 为 Message 类注册一个自定义反序列化器
        SimpleModule messageModule = new SimpleModule();
        messageModule.addDeserializer(Message.class, new MessageDeserializer());
        objectMapper.registerModule(messageModule);

        // register DeepResearchDeserializer， 为 OverAllState 类注册一个自定义反序列化器
        SimpleModule stateModule = new SimpleModule();
        stateModule.addDeserializer(OverAllState.class, new ItemQueryDeserializer(objectMapper));
        objectMapper.registerModule(stateModule);

        // other properties
        objectMapper.setVisibility(PropertyAccessor.FIELD, JsonAutoDetect.Visibility.ANY);	// 允许 ObjectMapper 访问和序列化类中的所有字段
        objectMapper.registerModule(new JavaTimeModule());
    }

    // 指定该序列化器处理的数据类型是 JSON 格式
    @Override
    public String contentType() {
        return "application/json";
    }

    @Override
    public void write(OverAllState object, ObjectOutput out) throws IOException {
        String json = objectMapper.writeValueAsString(object);
        byte[] jsonBytes = json.getBytes(StandardCharsets.UTF_8);
        out.writeInt(jsonBytes.length);
        out.write(jsonBytes);
    }

    @Override
    public OverAllState read(ObjectInput in) throws IOException{
        int length = in.readInt();
        byte[] jsonBytes = new byte[length];
        in.readFully(jsonBytes);
        String json = new String(jsonBytes, StandardCharsets.UTF_8);
        return objectMapper.readValue(json, OverAllState.class);
    }
}
